'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: retrieval_service.py
* @Time: 2025/11/15
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass
from uuid import UUID

from injector import inject
from langchain.retrievers import EnsembleRetriever
from langchain_core.documents import Document as LCDocument
from sqlalchemy.sql.expression import update

from pkg.sqlalchemy import SQLAlchemy

from .vector_database_service import VectorDataBaseService
from .jieba_service import JiebaService
from .base_service import BaseService
from ..entity.dataset_entity import RetrievalStrategy, RetrievalSource
from ..exception.exception import NotFoundException
from ..model import Dataset, DatasetQuery, Segment


@inject
@dataclass
class RetrievalService(BaseService):
    """检索服务"""
    db:SQLAlchemy
    jieba_service:JiebaService
    vector_database_service:VectorDataBaseService
    def search_in_datasets(self,
                           dataset_ids:list[UUID],
                           query:str,
                           retrieval_strategy:RetrievalStrategy.SEMANTIC,
                           k:int = 4,
                           score:float = 0.4,
                           retrieval_source:str=RetrievalSource.HIT_TESTING)->list[LCDocument]:
        """根据传递的query + 知识库列表执行检索， 并返回检索的文档 + 得分数据（如果是全文检索则得分为0）"""
        # todo:等待授权认证模块完成进行切换调整
        account_id = "e036e651-9eae-4b15-8dc3-b9d692121acf"
        #1、提取知识库列表并校验权限的同时更新知识库id
        datasets = self.db.session.query(Dataset).filter(
            Dataset.id.in_(dataset_ids),
            Dataset.account_id == account_id,
        ).all()
        if datasets is None or len(datasets) == 0:
            raise NotFoundException("该账户下无知识库库查询！")
        #2、构建不同的检索器类
        from internal.core.retrievers.semantic_retriever import SemanticRetriever
        from internal.core.retrievers.full_text_retriever import FullTextRetriever
        semantic_retriever = SemanticRetriever(
            dataset_ids=dataset_ids,
            vector_store=self.vector_database_service.vector_store,
            search_kwargs={
                "k":k,
                "score_threshold":score,
            }
        )

        full_text_retriever = FullTextRetriever(
            db=self.db,
            dataset_ids=dataset_ids,
            jieba_service=self.jieba_service,
            search_kwargs={
                "k":k
            }
        )

        hybrid_retrieval = EnsembleRetriever(
            retrievers=[semantic_retriever, full_text_retriever],
            weights=[0.5,0.5]
        )
        #3、根据不同的检索器执行检索
        if retrieval_strategy == RetrievalStrategy.SEMANTIC:
            lc_documents = semantic_retriever.invoke(query)[:k]
        elif retrieval_strategy == RetrievalStrategy.FULL_TEXT:
            lc_documents = full_text_retriever.invoke(query)[:k]
        else:
            lc_documents = hybrid_retrieval.invoke(query)[:k]
        #4、添加知识库查询记录，（只存储唯一记录， 也就是一个知识库如果检索了多篇文档， 也只存储一条）
        unique_dataset_ids = list(set(str(lc_document.metadata["dataset_id"]) for lc_document in lc_documents))
        for dataset_id in unique_dataset_ids:
            self.create(
                DatasetQuery,
                dataset_id=dataset_id,
                query=query,
                source=retrieval_source,
                #等待APP 完成后再配置
                source_app_id = None,
                created_by=account_id,
            )
        #5、批量更新片段的命中次数， 召回次数， 涵盖 构建 + 执行语句
        with self.db.auto_commit():
            stmt = (
                update(Segment).
                where(Segment.id.in_([lc_document.metadata["dataset_id"] for lc_document in lc_documents])).
                values(hit_count=Segment.hit_count + 1)
            )
        return lc_documents



